package com.portal.common.redis.service;

import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

@Slf4j
@Component
public class LimitAlgorithmService {
    /**
     * 固定窗口算法lua
     */
    public String gdckLuaScript() {
        StringBuilder lua = new StringBuilder();
        lua.append("local c");
        lua.append("\nc = redis.call('get',KEYS[1])");
        // 调用不超过最大值，则直接返回
        lua.append("\nif c and tonumber(c) > tonumber(ARGV[1]) then");
        lua.append("\nreturn c;");
        lua.append("\nend");
        // 执行计算器自加
        lua.append("\nc = redis.call('incr',KEYS[1])");
        lua.append("\nif tonumber(c) == 1 then");
        // 从第一次调用开始限流，设置对应键值的过期
        lua.append("\nredis.call('expire',KEYS[1],ARGV[2])");
        lua.append("\nend");
        lua.append("\nreturn c;");
        return lua.toString();
    }

    /**
     * 滑动窗口算法lua
     */
    public String hdckLuaScript() {
        StringBuilder sb = new StringBuilder();
        sb.append(" local key = KEYS[1] ");
        //sb.append(" -- 限流请求数 ");
        sb.append(" local limitCount = ARGV[1] ");
        //sb.append(" -- 限流开始时间戳（一般是当前时间减去前多少范围时间，例如前5秒） ");
        sb.append(" local startTime = ARGV[2] ");
        //sb.append(" -- 限流结束时间戳（当前时间） ");
        sb.append(" local endTime = ARGV[3] ");
        //sb.append(" -- 限流超时时间-用于清除内存-毫秒（默认与限制时间一致） ");
        sb.append(" local timeout = ARGV[4] ");
        //当前请求数
        sb.append(" local currentCount = redis.call('zcount', key, startTime, endTime)  ");
        //sb.append(" -- 限流存在并且超过限流大小，则返回剩余可用请求数=0 ");
        sb.append(" if (currentCount and tonumber(currentCount) >= tonumber(limitCount)) then ");
        sb.append("     return 0 ");
        sb.append(" end ");
        //sb.append(" -- 记录本次请求 ");
        sb.append(" redis.call('zadd', key, endTime, endTime) ");
        //sb.append(" -- 设置超时时间 ");
        sb.append(" redis.call('expire', key, timeout) ");
        //sb.append(" -- 返回剩余可用请求数 ");
        sb.append(" return tonumber(limitCount) - tonumber(currentCount) ");
        return sb.toString();
    }

    /**
     * 令牌桶算法lua
     */
    public String lptLuaScript(){
        StringBuilder sb = new StringBuilder();
        //sb.append(" --参数说明：key[1]为对应服务接口的信息，capacity为最大容量，rate为令牌生成速率（例如500ms生成一个则为0.5）,leftTokenNum为剩余令牌数,lastTs为时间戳 ");
        sb.append(" local limitInfo = redis.call('hmget', KEYS[1], 'capacity', 'rate','leftTokenNum', 'lastTs') ");
        sb.append(" local capacity = limitInfo[1] ");
        sb.append(" local rate = limitInfo[2] ");
        sb.append(" local leftTokenNum= limitInfo[3] ");
        sb.append(" local lastTs = limitInfo[4] ");
        // 本次需要令牌数
        sb.append(" local need = 1 ");
        //sb.append(" --初始化令牌桶 ");
        sb.append(" if capacity == false or rate == false or leftTokenNum == false then ");
        sb.append("     capacity = tonumber(ARGV[1]) ");
        sb.append("     rate = tonumber(ARGV[2]) ");
        sb.append("     leftTokenNum = tonumber(ARGV[1]) - need ");
        sb.append("     lastTs = tonumber(ARGV[3]) ");
        sb.append("     redis.call('hmset', KEYS[1], 'capacity', capacity, 'rate', rate, 'leftTokenNum', leftTokenNum, 'lastTs', lastTs) ");
        sb.append("     return leftTokenNum ");
        sb.append(" else ");
        sb.append(" 	local nowTs = tonumber(ARGV[3]) ");
        sb.append("     local createTokenNum = tonumber((nowTs - lastTs)* rate/1000) ");
        //sb.append(" 	--计算该段时间的剩余令牌（当前总令牌数） ");
        sb.append("     leftTokenNum = createTokenNum + leftTokenNum ");
        //sb.append(" 	--设置剩余令牌（留下最小数） ");
        sb.append("     leftTokenNum = math.min(capacity, leftTokenNum) ");
        //sb.append(" 	--设置本次请求的时间 ");
        sb.append("     lastTs = nowTs ");
        //sb.append("     --判断是否还有令牌 ");
        sb.append("     if leftTokenNum >= need then ");
        //sb.append("         --减去需要的令牌 ");
        sb.append("         leftTokenNum = leftTokenNum - need ");
        //sb.append("         --更新剩余空间和上一次的生成令牌时间戳 ");
        sb.append("         redis.call('hmset', KEYS[1], 'capacity', capacity, 'rate', rate,'leftTokenNum', leftTokenNum, 'lastTs', lastTs) ");
        sb.append("         return leftTokenNum ");
        sb.append("     end ");
        sb.append("     return -1 ");
        sb.append(" end ");
        return sb.toString();
    }

    /**
     * 漏桶算法lua
     */
    public String ltLuaScript(){
        StringBuilder sb = new StringBuilder();
        //sb.append(" --参数说明：key[1]为对应服务接口的信息，capacity为容量,passRate为漏水速率,addWater为每次请求加水量（默认为1）,water为当前水量,lastTs为时间戳 ");
        sb.append(" local limitInfo = redis.call('hmget', KEYS[1], 'capacity', 'passRate','water', 'lastTs') ");
        sb.append(" local capacity = limitInfo[1] ");
        sb.append(" local passRate = limitInfo[2] ");
        //加水量固定为1（一次请求）
        sb.append(" local addWater= 1 ");
        sb.append(" local water = limitInfo[3] ");
        sb.append(" local lastTs = limitInfo[4] ");
        //sb.append(" --初始化漏斗 ");
        sb.append(" if capacity == false or passRate == false then ");
        sb.append("     capacity = tonumber(ARGV[1]) ");
        sb.append("     passRate = tonumber(ARGV[2]) ");
        //sb.append("     --当前水量（第一次加水量） ");
        sb.append("     water = addWater ");
        sb.append("     lastTs = tonumber(ARGV[3]) ");
        sb.append("     redis.call('hmset', KEYS[1], 'capacity', capacity, 'passRate', passRate,'addWater',addWater,'water', water, 'lastTs', lastTs) ");
        sb.append("     return 1 ");
        sb.append(" else ");
        sb.append("     local nowTs = tonumber(ARGV[3]) ");
        //sb.append("     --计算距离上一次请求到现在的漏水量 ");
        sb.append("     local waterPass = tonumber((nowTs - lastTs)* passRate/1000) ");
        //sb.append("     --计算当前水量,即执行漏水 ");
        sb.append("     water=math.max(0,water-waterPass) ");
        //sb.append("     --设置本次请求的时间 ");
        sb.append("     lastTs = nowTs ");
        //sb.append("     --判断是否可以加水 ");
        sb.append("     addWater=tonumber(addWater) ");
        sb.append("     if capacity-water >= addWater then ");
        //sb.append("         --加水 ");
        sb.append("         water=water+addWater ");
        //sb.append("         --更新当前水量和时间戳 ");
        sb.append("         redis.call('hmset', KEYS[1], 'water', water, 'lastTs', lastTs) ");
        sb.append("         return 1 ");
        sb.append("     end ");
        sb.append("     return 0 ");
        sb.append(" end ");
        return sb.toString();
    }


}
